import time;
import math;

class Factor:
    def __init__(self, targetVar=None, aFactors=[]):
        self.targetVar = targetVar;
        self.aVariables = [];
        self.mVar2Vals = dict();
        self.aFactors = [];
        self.aMax, self.aMaxVal, self.bCached = [], [], False;
        self.iMax = -1;
        for f in aFactors:
            self.addFactor(f); 
        
    def addFactor(self, factor):
        self.aFactors.append(factor);
        self.addVariables(factor.aVariables);

    def addVariables(self, aVariables):
        for var in aVariables:
            if not(var in self.aVariables) and var!=self.targetVar:
                self.aVariables.append(var);
                self.mVar2Vals[var] = 0;
                
    def addTargetVar(self, targetVar):
        self.targetVar = targetVar;
        
    def createCache(self):
        nVars = len(self.aVariables);
#         nTime = 0;
        for i in range(2**nVars): #For Each Combination of values, get their iMax & maxVal
            sBin = str(bin(i)[2:]).zfill(nVars);
            for k,var in enumerate(self.aVariables):
                self.mVar2Vals[var] = int(sBin[k]);
#             nInitT = time.time();
            self.aMaxVal.append(self.getValue(self.mVar2Vals));     
#             nTime += (time.time()-nInitT);
#             print i, (time.time()-nInitT);
            self.aMax.append(self.iMax);
#         if nTime/2**nVars>0:
#             print "AvgTime:", nTime/2**nVars;
        self.bCached = True;
    
    def getVariableNames(self):
        aNames = [];
        for v in self.aVariables:
            aNames.append(v.name);
        return aNames;
    
    def getValue(self, mVar2Vals): #, nLevel=0):
        if self.bCached:
            return self.getCachedValue(mVar2Vals);
        nMax = float("-inf");
        for i in range(2):
            nFactor = 0;
            mVar2Vals[self.targetVar] = i;
            for f in self.aFactors:
                mVar2Vals_F = dict();
                for var in f.aVariables:
                    mVar2Vals_F[var] = mVar2Vals[var];
                mVar2Vals_F[self.targetVar] = i;
#                 for kVar in mVar2Vals:
#                     if f.mVar2Vals.has_key(kVar):
#                         mVar2Vals_F[kVar] = mVar2Vals[kVar];
                nF = f.getValue(mVar2Vals_F); #, nLevel+1);
                nFactor += nF;
            if nFactor>=nMax:
                nMax = nFactor;
                self.iMax = i;
        return nMax;
    
    def getCachedValue(self, mVar2Vals):
        sBin = "";
        for var in self.aVariables:
            if not(mVar2Vals.has_key(var)):
                print "Error: Couldn't find CPT Variable:", var.name, "for target variable:", self.targetVar.name;
            sVal = str(mVar2Vals[var]) if mVar2Vals.has_key(var) else "0";
            sBin += sVal;
        i = int(sBin,2) if sBin!="" else 0;
        self.iMax = self.aMax[i]
        nMax = self.aMaxVal[i];
        return nMax;
        
    
    
class FactorCPT (Factor):
    def __init__(self, aVariables, aCPT, bEv=False):
        self.aVariables = aVariables; #1st item is "targetVar" (only for FactorCPTs)
        self.mVar2Vals = dict();
        for var in aVariables:
            self.mVar2Vals[var] = 0;
        self.aCPT = aCPT;
        self.bEvidence = bEv;
    
    def getValue(self, mVar2Vals): #, nLevel=0):
        sBin = ""; # if len(mVar2Vals)>1 else "0";
        for var in self.aVariables:
            sBin += str(mVar2Vals[var]);
        sBin = "1"+sBin if self.bEvidence else sBin;
        nQueryVal = int(sBin[0]);
        #If it's Prior P(X), then the index in CPT is always "0" 
        #(and nQueryVal takes care of whether to get P(X=False) or P(X=True))
        sBin = sBin[1:] if len(mVar2Vals)>1 else "0";
        iP = int(sBin,2) if sBin!="" else 0;
        nProb = 1-self.aCPT[iP] if nQueryVal==0 else self.aCPT[iP];
        logProb = float("-inf") if nProb==0 else math.log(nProb);
        return logProb;
    

class VarElimination(object):
    def __init__(self):
        pass;
    
    def getInference(self, g, aNodes, aObsNodes):
        aFactorsCPT = [];        
        aNodes2 = [];
        for node in aNodes:
            if len(node.aSubjects)==1:
                aNodes2.append(node);
        for node in aNodes:
            if len(node.aSubjects)==2:
                aNodes2.append(node);
        
        for node in aNodes:
            aFactorsCPT.append(FactorCPT([node] + g.ndict[node][0], g.ndict[node][1]));
        for obsN in aObsNodes:
            aFactorsCPT.append(FactorCPT(g.ndict[obsN][0], g.ndict[obsN][1], True));
#         tick = time.time();
        return self.getMaxProd_VE(aFactorsCPT, aNodes2);
#         print time.time()-tick, "seconds";
     
    def getMaxProd_VE(self, aFactors, aVariables):
        aNewF = [];
        nMaxArgs = -1;
        nArgs, nOps = 0, 0;
        print "Eliminating Vars... ("+str(len(aVariables))+" nodes)";
        nInitT = time.time();
        for i, var in enumerate(aVariables):
            aFactors = self.eliminateVar(aFactors, var);
            aNewF.append(aFactors[-1]);
            aVars = [];
            for varF in aFactors[-1].aVariables:
                aVars.append(varF.name);
            print var.name+":", len(aFactors[-1].aVariables), "-", aVars;
            if len(aFactors[-1].aVariables)>nMaxArgs:
                nMaxArgs = len(aFactors[-1].aVariables);
#             print var.name, len(aFactors[-1].aVariables);
            nArgs += len(aFactors[-1].aVariables);
            nOps += 2**len(aFactors[-1].aVariables);
            if (i+1)%100==0: print str(i+1)+ "/" +str(len(aVariables))+" Nodes";
#         print "---\n"
        print "nArgs:", nArgs, "nOps:", nOps;
        print "Finished with nMaxArgs:", nMaxArgs, "in", time.time()-nInitT, "seconds.";
        print "TraceBack...";
        return self.traceBack(aNewF);
    
    def eliminateVar(self, aFactors, variable):
        newF = Factor();
        newF.addTargetVar(variable);
        #Add Factors
        aDelFs = [];
        for i,f in enumerate(aFactors):
            if variable in f.aVariables:
                newF.addFactor(f);
                newF.addVariables(f.aVariables);
                aDelFs.append(i);
#         nInitT = time.time();
        newF.createCache();
#         print "CreateCache with", len(newF.aVariables), "variables: %2.4f" % (time.time()-nInitT), "seconds. Node:", newF.targetVar.name
#         if len(newF.aVariables)>=6:
#             print "CreateCache with", len(newF.aVariables), "variables:", time.time()-nInitT, "seconds. Node:", newF.targetVar.name;
        #Delete considered Factors
        aDelFs.reverse();
        for i in aDelFs:
            del aFactors[i];
        #Add Created Factor
        aFactors.append(newF);
        return aFactors;
    
    def traceBack(self, aFactors):
        aVariables, aValues, aFactorVals = [], [], [];
        mVar2Vals = dict();
        nInitT = time.time();
        aFactors.reverse();
#         print "---\n"
        for f in aFactors:
#             print f.targetVar.name;
            nF = f.getValue(mVar2Vals);
            mVar2Vals[f.targetVar] = f.iMax;
            aVariables.append(f.targetVar);
            aValues.append(f.iMax);
            aFactorVals.append(nF);
#         print "---\n"
#         for i,var in enumerate(aVariables):
#             print var.name, "=", aValues[i], aFactorVals[i];
        print "TraceBack Finished", "in", time.time()-nInitT, "seconds.";
        return [aVariables, aValues];
